import time
from unitree_sdk2py.core.channel import ChannelFactoryInitialize
from unitree_sdk2py.go2.sport.sport_client import SportClient


class Go2Controller:
    """
    宇树Go2 运动控制封装类
    适用于：Unitree SDK2 Python接口
    """

    def __init__(self, interface="enx00e0986113a6", timeout=10.0, auto_init=True):
        """
        初始化控制器
        :param interface: 网卡接口名称（如 enx00e0986113a6）
        :param timeout: 超时时间（秒）
        :param auto_init: 是否自动初始化连接
        """
        self.interface = interface
        self.timeout = timeout
        self.sport_client = None
        if auto_init:
            self.init_connection()

    def init_connection(self):
        """初始化与Go2的连接"""
        try:
            print(f"🟡 初始化Go2连接接口: {self.interface}")
            ChannelFactoryInitialize(0, self.interface)
            self.sport_client = SportClient()
            self.sport_client.SetTimeout(self.timeout)
            self.sport_client.Init()
            print("✅ Go2控制器初始化完成")
        except Exception as e:
            print(f"❌ 初始化失败: {e}")
            self.sport_client = None

    def _safe_call(self, func_name, wait_time=0.3):
        """线程安全的动作执行"""
        if not self.sport_client:
            print("❌ SportClient 未初始化，请检查连接")
            return False
            
        try:
            # 异步执行动作
            import threading
            def _async_exec():
                try:
                    func = getattr(self.sport_client, func_name)
                    func()
                    print(f"✅ 执行动作: {func_name}（等待 {wait_time:.1f}s）")
                    
                    # 非阻塞延迟
                    start_time = time.time()
                    while time.time() - start_time < wait_time:
                        time.sleep(0.05)  # 更细粒度的检查
                        if not self.sport_client:
                            return False
                except Exception as e:
                    print(f"❌ 执行动作失败: {func_name}, 错误: {e}")
            
            thread = threading.Thread(target=_async_exec)
            thread.daemon = True  # 设置为守护线程
            thread.start()
            return True
        except Exception as e:
            print(f"❌ 执行动作失败: {func_name}, 错误: {e}")
            return False

    # -------------------- 基础动作 --------------------
    def stand_up(self):
        """站立"""
        return self._safe_call("StandUp", wait_time=4.0)

    def stand_down(self):
        """下蹲"""
        return self._safe_call("StandDown", wait_time=4.0)

    def damp(self):
        """阻尼模式"""
        return self._safe_call("Damp", wait_time=2.0)

    def walk_ready(self):
        """行走准备姿态"""
        return self._safe_call("WalkReady", wait_time=2.0)

    def balance_stand(self):
        """平衡站立"""
        return self._safe_call("BalanceStand", wait_time=3.0)

    def stand_high(self):
        """抬高站姿"""
        return self._safe_call("StandHigh", wait_time=3.0)

    def stretch(self):
        """拉伸动作"""
        return self._safe_call("Stretch", wait_time=3.0)

    def hand_shake(self):
        """挥手（部分Go2版本支持）"""
        return self._safe_call("HandShake", wait_time=3.0)

    # -------------------- 组合动作 --------------------
    def sit(self):
        """从站立 → 下蹲"""
        print("🟢 执行：坐下（站立→下蹲）")
        self.stand_up()
        self.stand_down()

    # -------------------- 表演动作 --------------------
    def happy(self):
        """开心"""
        return self._safe_call("Content", wait_time=3.0)

    def wallow(self):
        """打滚"""
        return self._safe_call("Wallow", wait_time=4.0)

    def pose(self):
        """摆姿势"""
        return self._safe_call("Pose", wait_time=3.0)

    def scrape(self):
        """拜年"""
        return self._safe_call("Scrape", wait_time=5.0)

    def wake_up(self):
        """从蹲下 → 站起"""
        print("🟢 执行：起立（下蹲→站立）")
        self.stand_down()
        self.stand_up()

    def relax(self):
        """阻尼模式+下蹲"""
        print("🟢 执行：放松模式")
        self.damp()
        self.stand_down()

    def demo_pose(self):
        """演示模式"""
        print("🟢 执行：演示姿态")
        self.stand_up()
        self.stretch()
        self.stand_down()

    # -------------------- 舞蹈动作 --------------------
    def Dance1(self):
        """
        表演动作1
        功能：调用 Go2 的 Dance1 接口
        """
        return self._safe_call("Dance1", wait_time=5.0)

    def Dance2(self):
        """
        表演动作2
        功能：调用 Go2 的 Dance2 接口
        """
        return self._safe_call("Dance2", wait_time=5.0)

    # -------------------- 安全控制 --------------------
    def shutdown(self):
        """安全关闭连接"""
        print("🛑 正在关闭Go2连接")
        self.sport_client = None
        
    def close(self):
        """关闭控制器并释放资源"""
        self.shutdown()
        # 强制终止所有后台线程
        import threading
        for thread in threading.enumerate():
            if thread != threading.current_thread():
                thread.join(timeout=0.1)  # 等待线程终止（最多100ms）


# -------------------- 测试 --------------------
if __name__ == "__main__":
    go2 = Go2Controller(interface="enx00e0986113a6")

    go2.stand_down()
    go2.stand_up()
    go2.damp()
    go2.walk_ready()
    go2.shutdown()
